package top.sqdpt.rag.transformer;

import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.rag.query.transformer.CompressingQueryTransformer;
import dev.langchain4j.rag.query.transformer.ExpandingQueryTransformer;
import org.springframework.util.CollectionUtils;

import java.util.Arrays;
import java.util.Collection;

public class CompressingAndExpandingQueryTransformer extends CompressingQueryTransformer {

    public CompressingAndExpandingQueryTransformer(ChatLanguageModel chatLanguageModel) {
        super(chatLanguageModel);
    }

    @Override
    public Collection<Query> transform(Query query) {
        Collection<Query> transform = super.transform(query);
        if (CollectionUtils.isEmpty(transform)) {
            return transform;
        }
        Query compressingQuery = transform.stream().findFirst().orElse(null);
        ExpandingQueryTransformer expandingQueryTransformer = new ExpandingQueryTransformer(chatLanguageModel);
        return expandingQueryTransformer.transform(compressingQuery);
    }
}
